import os
import json
import chardet
from mutagen import File
from mutagen.flac import FLAC
from mutagen.mp3 import MP3
from mutagen.mp4 import MP4
from mutagen.oggvorbis import OggVorbis
from mutagen.oggopus import OggOpus
from mutagen.wave import WAVE
from mutagen.apev2 import APEv2File
from tqdm import tqdm

SUPPORTED_FORMATS = {
    '.mp3': MP3,
    '.flac': FLAC,
    '.m4a': MP4,
    '.ogg': OggVorbis,
    '.opus': OggOpus,
    '.wav': WAVE,
    '.ape': APEv2File,
    '.wv': APEv2File  # WavPack
}

def detect_encoding(file_path):
    """检测文件编码"""
    with open(file_path, 'rb') as f:
        result = chardet.detect(f.read())
    return result['encoding'] or 'utf-8'

def extract_metadata(file_path):
    """提取音频文件元数据"""
    ext = os.path.splitext(file_path)[1].lower()
    if ext not in SUPPORTED_FORMATS:
        return None

    try:
        audio = File(file_path, easy=True)
        if audio is None:
            return None

        # 基础技术信息
        info = {
            'format': ext[1:].upper(),
            'file_size': os.path.getsize(file_path),
            'duration': audio.info.length,
            'bitrate': getattr(audio.info, 'bitrate', 0),
            'sample_rate': getattr(audio.info, 'sample_rate', 0),
            'channels': getattr(audio.info, 'channels', 0),
        }

        # 标签元数据
        tags = {}
        if hasattr(audio, 'tags') and audio.tags is not None:
            for key, value in audio.tags.items():
                # 处理多值标签和文本编码
                clean_value = [v.decode('utf-8', 'ignore') if isinstance(v, bytes) else v for v in value]
                tags[key] = clean_value[0] if len(clean_value) == 1 else clean_value

        # 标准化常见标签名
        metadata = {
            'technical': info,
            'tags': {
                'title': tags.get('title', [''])[0],
                'artist': tags.get('artist', [''])[0],
                'album': tags.get('album', [''])[0],
                'track_number': tags.get('tracknumber', ['0/0'])[0].split('/')[0],
                'genre': tags.get('genre', [''])[0],
                'date': tags.get('date', [''])[0],
                'album_artist': tags.get('albumartist', [''])[0],
                'composer': tags.get('composer', [''])[0],
            }
        }

        # 添加原始标签供参考
        metadata['raw_tags'] = tags

        return metadata

    except Exception as e:
        print(f"Error processing {file_path}: {str(e)}")
        return None

def process_audio_files(input_dir, output_dir):
    """处理所有音频文件"""
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)

    files = [f for f in os.listdir(input_dir) 
             if os.path.splitext(f)[1].lower() in SUPPORTED_FORMATS]

    for filename in tqdm(files, desc="Processing files"):
        input_path = os.path.join(input_dir, filename)
        output_path = os.path.join(output_dir, f"{os.path.splitext(filename)[0]}.json")

        metadata = extract_metadata(input_path)
        if metadata:
            with open(output_path, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, ensure_ascii=False, indent=2)

if __name__ == "__main__":
    INPUT_DIR = "input"
    OUTPUT_DIR = "temp"

    print("Starting metadata extraction...")
    process_audio_files(INPUT_DIR, OUTPUT_DIR)
    print("\nMetadata extraction completed!")